#include "gmsmx.h"
#include <stdlib.h>
#include <sys/time.h>
#include <unistd.h>

int ec_key_print(EC_KEY* key);
int ec_key_gen_test();
int external_ec_key_test();
int sm2_enc_and_dec_test();
int sm2_key_exchang_test();
int sm2_sign_and_verify_test();
int sm3_hash_test();
int sm4_cbc_test();
int sm4_ecb_test();
int sm2_enc_performance_test();
int sm2_dec_performance_test();
int sm2_sign_performance_test();
int sm2_verify_performance_test();
int sm3_performance_test();
int sm4_enc_cbc_performance_test();
int sm4_dec_cbc_performance_test();
int sm4_enc_ecb_performance_test();
int sm4_dec_ecb_performance_test();
int test_all();

typedef struct test_item {
    char* des;
    int (*func)(void);
} test_item;

struct test_item item[] = {
    {"ecc key pair generation test", ec_key_gen_test},
    {"externel ecc keypair test", external_ec_key_test},
    {"SM2 encrypt and decrypt test", sm2_enc_and_dec_test},
    {"SM2 sign and verify test", sm2_sign_and_verify_test},
    {"SM2 key exchange test", sm2_key_exchang_test},
    {"SM3 hash test", sm3_hash_test},
    {"SM4 cbc test(not need padding)", sm4_cbc_test},
    {"SM4 ecb test", sm4_ecb_test},
    {"SM2 enc performance test", sm2_enc_performance_test},
    {"SM2 dec performance test", sm2_dec_performance_test},
    {"SM2 sign performance test", sm2_sign_performance_test},
    {"SM2 verify performance test", sm2_verify_performance_test},
    {"SM3 hash performance test", sm3_performance_test},
    {"SM4 encrypt[cbc] performance test", sm4_enc_cbc_performance_test},
    {"SM4 decrypt[cbc] performance test", sm4_dec_cbc_performance_test},
    {"SM4 encrypt[ecb] performance test", sm4_enc_ecb_performance_test},
    {"SM4 decrypt[ecb] performance test", sm4_dec_ecb_performance_test},
    {"test all", test_all},
};

int atoi_c(const char* buf, int* result)
{
    char* invalid;
    *result = (int)strtol(buf, &invalid, 10);
    if (*invalid != '\0') return -1;
    return 0;
}

void help_print(void)
{
    int i;
    int num = sizeof(item) / sizeof(item[0]);

    printf(
        "\n----------------------"
        "\nselect test item:\n");
    for (i = 0; i < num; i++) {
        printf("\t%02d : %s\n", i+1, item[i].des);
    }
    printf("\tother : exit\n");
    printf("select : ");
}

int main(int argc, char const* argv[])
{
    int  rc = 0;
    int  input;
    char inputs[1024];
    int  num = sizeof(item) / sizeof(item[0]);

    do {
        help_print();
        scanf("%s", inputs);
        rc = atoi_c(inputs, &input);
        if (rc < 0 || input > num || input < 1) {
            break;
        }
        printf("----------------------\n");
        item[input-1].func();
    } while (1);

    return 0;
}

int ec_key_print(EC_KEY* key)
{
    unsigned char pub[64]  = {0};
    unsigned char priv[32] = {0};
    ec_key_to_buf(key, pub, priv);
    printf("EC_KEY:\npubkey:");
    print_hex(pub, 64);
    printf("\nprivkey:");
    print_hex(priv, 32);
    return 0;
}

int ec_key_gen_test()
{
    EC_KEY* key = ec_key_new();
    ec_gen_keypair(key);

    int ret = ec_check_pubkey(key);
    if (ret) {
        printf("pubkey check failed. %d\n", ret);
    } else {
        printf("pubkey check ok\n");
    }

    ec_key_free(key);
    return 0;
}

int sm2_sign_and_verify_test()
{
    const char*   msg = "message digst";
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;

    EC_KEY* key = ec_key_new();
    ec_gen_keypair(key);

    int ret = ec_check_pubkey(key);
    if (ret) {
        printf("pubkey check failed. %d\n", ret);
    } else {
        printf("pubkey check ok\n");
    }

    SM2_compute_message_digest((unsigned char*)msg, strlen(msg),
                               (unsigned char*)id, strlen(id), dgst, key);
    // print_hex(dgst, SM3_DIGEST_LENGTH);

    ret = SM2_sign(dgst, &rs, key);
    if (ret) {
        printf("sm2 sign failed. %d\n", ret);
    } else {
        printf("sm2 sign ok\n");
    }

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_verify(dgst, &rs, key);
    if (ret) {
        printf("sm2 verify failed. %d\n", ret);
    } else {
        printf("sm2 verify ok\n");
    }
    ec_key_free(key);
    return 0;
}

unsigned char private_1[32] = {0x39, 0x45, 0x20, 0x8f, 0x7b, 0x21, 0x44, 0xb1,
                               0x3f, 0x36, 0xe3, 0x8a, 0xc6, 0xd3, 0x9f, 0x95,
                               0x88, 0x93, 0x93, 0x69, 0x28, 0x60, 0xb5, 0x1a,
                               0x42, 0xfb, 0x81, 0xef, 0x4d, 0xf7, 0xc5, 0xb8};

unsigned char public_1[64] = {
    0x09, 0xF9, 0xDF, 0x31, 0x1E, 0x54, 0x21, 0xA1, 0x50, 0xDD, 0x7D,
    0x16, 0x1E, 0x4B, 0xC5, 0xC6, 0x72, 0x17, 0x9F, 0xAD, 0x18, 0x33,
    0xFC, 0x07, 0x6B, 0xB0, 0x8F, 0xF3, 0x56, 0xF3, 0x50, 0x20, 0xCC,
    0xEA, 0x49, 0x0C, 0xE2, 0x67, 0x75, 0xA5, 0x2D, 0xC6, 0xEA, 0x71,
    0x8C, 0xC1, 0xAA, 0x60, 0x0A, 0xED, 0x05, 0xFB, 0xF3, 0x5E, 0x08,
    0x4A, 0x66, 0x32, 0xF6, 0x07, 0x2D, 0xA9, 0xAD, 0x13,
};

int external_ec_key_test()
{
    EC_KEY key;
    key.pub    = mpi_point_new(0);
    key.priv   = mpi_read_raw_data(private_1, 32);
    key.pub->x = mpi_read_raw_data(public_1, 32);
    key.pub->y = mpi_read_raw_data(public_1 + 32, 32);
    key.pub->z = mpi_new(0);
    mpi_set_ui(key.pub->z, 1);

    // ec_key_print(&key);
    int ret = ec_check_pubkey(&key);
    if (ret) {
        printf("pubkey check failed. %d\n", ret);
    } else {
        printf("pubkey check ok\n");
    }
    const char*   msg = "message digst";
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;

    SM2_compute_message_digest((unsigned char*)msg, strlen(msg),
                               (unsigned char*)id, strlen(id), dgst, &key);
    // print_hex(dgst, SM3_DIGEST_LENGTH);

    ret = SM2_sign(dgst, &rs, &key);
    if (ret) {
        printf("sm2 sign failed. %d\n", ret);
    } else {
        printf("sm2 sign ok\n");
    }

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_verify(dgst, &rs, &key);
    if (ret) {
        printf("sm2 verify failed. %d\n", ret);
    } else {
        printf("sm2 verify ok\n");
    }
    return 0;
}

int sm2_enc_and_dec_test()
{
    const char*   in = "encryption standard";
    unsigned char cipher[128];
    unsigned char out[128] = {0};
    size_t        outlen;
    int           ret;

    EC_KEY* key = ec_key_new();
    ec_gen_keypair(key);

    printf("encrypt int:%s, len=%ld\n", in, strlen(in));

    ret =
        SM2_encrypt((const unsigned char*)in, strlen(in), cipher, &outlen, key);
    if (ret) {
        printf("sm2 encrypt failed. %d\n", ret);
        goto free;
    } else {
        printf("sm2 encrypt ok\n");
    }
    printf("cipher(%ld):", outlen);
    print_hex(cipher, outlen);

    // print_hex(dgst, SM3_DIGEST_LENGTH);
    ret = SM2_decrypt(cipher, outlen, out, &outlen, key);
    if (ret) {
        printf("sm2 decrypt failed. %d\n", ret);
        goto free;
    } else {
        printf("sm2 decrypt ok\n");
    }
    printf("decrypt out:%s, len=%ld\n", out, outlen);
free:
    ec_key_free(key);
    return 0;
}

#define KEYLEN 16
int sm2_key_exchang_test()
{
    int         rc  = 0;
    const char* Aid = "123456";
    const char* Bid = "123456";

    unsigned char Ra[SM2_PUBKEY_LEN];
    unsigned char Rb[SM2_PUBKEY_LEN];
    unsigned char kab[KEYLEN];
    unsigned char kba[KEYLEN];
    unsigned char suma[SM3_DIGEST_LENGTH];
    unsigned char sumb[SM3_DIGEST_LENGTH];

    EC_KEY* keyA = ec_key_new();
    EC_KEY* keyB = ec_key_new();
    ec_gen_keypair(keyA);
    ec_gen_keypair(keyB);

    SM2_KAP_CTX ctxA;
    SM2_KAP_CTX ctxB;

    rc = SM2_KAP_CTX_init(&ctxA, keyA, keyB->pub, Aid, strlen(Aid), Bid,
                          strlen(Bid), 1, 1);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_CTX_init(&ctxB, keyB, keyA->pub, Bid, strlen(Bid), Aid,
                          strlen(Aid), 0, 1);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_prepare(&ctxA, Ra);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_prepare(&ctxB, Rb);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_compute_key(&ctxA, Rb, kab, KEYLEN, suma);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }
    rc = SM2_KAP_compute_key(&ctxB, Ra, kba, KEYLEN, sumb);
    if (rc) {
        printf("error rc = %d\n", rc);
        goto end;
    }

    rc = SM2_KAP_final_check(&ctxA, sumb);
    if (rc) {
        printf("error rc = %d\n", rc);
    }
    rc = SM2_KAP_final_check(&ctxB, suma);
    if (rc) {
        printf("error rc = %d\n", rc);
    }

    printf("key:\n");
    print_hex(kab, KEYLEN);
    print_hex(kba, KEYLEN);
end:
    SM2_KAP_CTX_cleanup(&ctxA);
    SM2_KAP_CTX_cleanup(&ctxB);
    ec_key_free(keyA);
    ec_key_free(keyB);
    return 0;
}

int sm3_hash_test()
{
    unsigned char in[] = {0x61, 0x62, 0x63};
    unsigned char dgst[SM3_DIGEST_LENGTH];

    printf("input data:");
    print_hex(in, sizeof(in));
    sm3(in, sizeof(in), dgst);

    printf("output data:");
    print_hex(dgst, SM3_DIGEST_LENGTH);
    return 0;
}

int sm4_cbc_test()
{
    #define DATALEN 13
    unsigned char in[DATALEN];
    unsigned char out[SM4_BLOCK_SIZE];

    unsigned char ukey[SM4_KEY_LENGTH];
    unsigned char iv[SM4_IV_LENGTH];
    unsigned char iv_d[SM4_IV_LENGTH];

    random_generator(in, DATALEN);
    random_generator(ukey, SM4_KEY_LENGTH);
    random_generator(iv, SM4_IV_LENGTH);
    memcpy(iv_d, iv, SM4_IV_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);
    printf("plain data:");
    print_hex(in, DATALEN);

    sm4_cbc_encrypt(in, out, DATALEN, &key, iv, 1);
    printf("cipher data:");
    print_hex(out, SM4_BLOCK_SIZE);

    sm4_set_decrypt_key(&key, ukey);
    sm4_cbc_encrypt(out, in, SM4_BLOCK_SIZE, &key, iv_d, 0);
    printf("decrypt data:");
    print_hex(in, DATALEN);
    return 0;
}

int sm4_ecb_test()
{
    unsigned char in[SM4_BLOCK_SIZE];
    unsigned char out[SM4_BLOCK_SIZE];
    unsigned char ukey[SM4_KEY_LENGTH];

    random_generator(in, SM4_BLOCK_SIZE);
    random_generator(ukey, SM4_KEY_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);
    printf("plain data:");
    print_hex(in, SM4_BLOCK_SIZE);

    sm4_ecb_encrypt(in, out, &key, 1);
    printf("cipher data:");
    print_hex(out, SM4_BLOCK_SIZE);

    sm4_set_decrypt_key(&key, ukey);
    sm4_ecb_encrypt(in, out, &key, 0);
    printf("decrypt data:");
    print_hex(in, SM4_BLOCK_SIZE);
    return 0;
}

static long get_millis(struct timeval* start, struct timeval* end)
{
    time_t      seconds = end->tv_sec - start->tv_sec;
    suseconds_t ms      = end->tv_usec - start->tv_usec;
    return seconds * 1000 + ms / 1000;
}

#define LOOP     1000
#define DATA_ELN 1024

int sm2_enc_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP;
    EC_KEY* key = ec_key_new();
    int ret;
    unsigned char in[DATA_ELN];
    unsigned char cipher[DATA_ELN+128];
    size_t outlen;
    ec_gen_keypair(key);
    random_generator(in, DATA_ELN);
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        ret = SM2_encrypt(in, DATA_ELN, cipher, &outlen, key);
        if (ret) {
            printf("sm2 encrypt failed. %d\n", ret);
            goto free;
        }  
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\n\n", diff, avg,
           1000 / avg);
free:
    ec_key_free(key);
    return 0;
}

int sm2_dec_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP;
    EC_KEY* key = ec_key_new();
    int ret;
    unsigned char in[DATA_ELN];
    unsigned char cipher[DATA_ELN+128];
    unsigned char out[DATA_ELN];
    size_t outlen;
    ec_gen_keypair(key);
    random_generator(in, DATA_ELN);
    ret = SM2_encrypt(in, DATA_ELN, cipher, &outlen, key);
    if (ret) {
        printf("sm2 encrypt failed. %d\n", ret);
    }  
    size_t inlen = outlen;
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        ret = SM2_decrypt(cipher, inlen, out, &outlen, key);
        if (ret) {
            printf("sm2 decrypt failed. %d\n", ret);
            goto free;
        }
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\n\n", diff, avg,
           1000 / avg);
free:
    ec_key_free(key);
    return 0;
}

int sm2_sign_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP;
    EC_KEY* key = ec_key_new();
    int ret;
    unsigned char in[DATA_ELN];
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;
    ec_gen_keypair(key);
    random_generator(in, DATA_ELN);
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        SM2_compute_message_digest(in, DATA_ELN,
            (unsigned char*)id, strlen(id), dgst, key);
        ret = SM2_sign(dgst, &rs, key);
        if (ret) {
            printf("sm2 sign failed. %d\n", ret);
            goto free;
        }
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\n\n", diff, avg,
           1000 / avg);
free:
    ec_key_free(key);
    return 0;
}

int sm2_verify_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP;
    EC_KEY* key = ec_key_new();
    int ret;
    unsigned char in[DATA_ELN];
    const char*   id  = "614837785@qq.com";
    unsigned char dgst[SM3_DIGEST_LENGTH];
    SM2_SIG_CTX   rs;
    ec_gen_keypair(key);
    random_generator(in, DATA_ELN);
    SM2_compute_message_digest(in, DATA_ELN,
        (unsigned char*)id, strlen(id), dgst, key);
    ret = SM2_sign(dgst, &rs, key);
    if (ret) {
        printf("sm2 sign failed. %d\n", ret);
        goto free;
    }
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        ret = SM2_verify(dgst, &rs, key);
        if (ret) {
            printf("sm2 verify failed. %d\n", ret);
            goto free;
        }
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\n\n", diff, avg,
           1000 / avg);
free:
    ec_key_free(key);
    return 0;
}

int sm3_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP*100;
    unsigned char in[DATA_ELN];
    unsigned char dgst[SM3_DIGEST_LENGTH];
    random_generator(in, DATA_ELN);
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        sm3(in, DATA_ELN, dgst);
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\tMbps=%f\n\n",
        diff, avg, 1000 / avg, 1000 / avg * DATA_ELN * 8 / 1000000);
    return 0;
}

int sm4_enc_cbc_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP*100;
    unsigned char in[DATA_ELN];
    unsigned char out[DATA_ELN];
    unsigned char ukey[SM4_KEY_LENGTH];
    unsigned char iv[SM4_IV_LENGTH];

    random_generator(in, DATA_ELN);
    random_generator(ukey, SM4_KEY_LENGTH);
    random_generator(iv, SM4_IV_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);

    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        sm4_cbc_encrypt(in, out, DATA_ELN, &key, iv, 1);
    }

    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\tMbps=%f\n\n",
        diff, avg, 1000 / avg, 1000 / avg * DATA_ELN * 8 / 1000000);
    return 0;
}

int sm4_dec_cbc_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP*100;
    unsigned char in[DATA_ELN];
    unsigned char out[DATA_ELN];
    unsigned char ukey[SM4_KEY_LENGTH];
    unsigned char iv[SM4_IV_LENGTH];

    random_generator(in, DATA_ELN);
    random_generator(ukey, SM4_KEY_LENGTH);
    random_generator(iv, SM4_IV_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);
    sm4_cbc_encrypt(in, out, DATA_ELN, &key, iv, 1);
    sm4_set_decrypt_key(&key, ukey);
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        sm4_cbc_encrypt(out, in, DATA_ELN, &key, iv, 0);
    }

    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\tMbps=%f\n\n",
        diff, avg, 1000 / avg, 1000 / avg * DATA_ELN * 8 / 1000000);
    return 0;
}

int sm4_enc_ecb_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP*100;
    unsigned char in[DATA_ELN];
    unsigned char out[DATA_ELN];
    unsigned char ukey[SM4_KEY_LENGTH];
    unsigned char iv[SM4_IV_LENGTH];

    random_generator(in, DATA_ELN);
    random_generator(ukey, SM4_KEY_LENGTH);
    random_generator(iv, SM4_IV_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);

    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        for (size_t i = 0; i < DATA_ELN; i += 16)
            sm4_ecb_encrypt(&in[i], &out[i], &key, 1);
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\tMbps=%f\n\n",
        diff, avg, 1000 / avg, 1000 / avg * DATA_ELN * 8 / 1000000);
    return 0;
    return 0;
}

int sm4_dec_ecb_performance_test()
{
    struct timeval start, end;
    size_t loop = LOOP*100;
    unsigned char in[DATA_ELN];
    unsigned char out[DATA_ELN];
    unsigned char ukey[SM4_KEY_LENGTH];
    unsigned char iv[SM4_IV_LENGTH];

    random_generator(in, DATA_ELN);
    random_generator(ukey, SM4_KEY_LENGTH);
    random_generator(iv, SM4_IV_LENGTH);

    sm4_key_t key;
    sm4_set_encrypt_key(&key, ukey);
    for (size_t i = 0; i < DATA_ELN; i += 16)
        sm4_ecb_encrypt(&in[i], &out[i], &key, 1);
    sm4_set_decrypt_key(&key, ukey);
    printf("test start , data len = %d\n", DATA_ELN);
    gettimeofday(&start, NULL);
    for (size_t i = 0; i < loop; i++)
    {
        for (size_t i = 0; i < DATA_ELN; i += 16)
            sm4_ecb_encrypt(&in[i], &out[i], &key, 0);
    }
    gettimeofday(&end, NULL);
    double diff = get_millis(&start, &end);
    double avg  = diff / loop;
    printf("finished. test result(ms): total=%f\tavg=%f\tpps=%f\tMbps=%f\n\n",
        diff, avg, 1000 / avg, 1000 / avg * DATA_ELN * 8 / 1000000);
    return 0;
}

int test_all()
{
    int num = sizeof(item) / sizeof(item[0]);
    for (int i = 0; i < num - 1; i++) {
        printf("test \t%02d : %s ...... \n", i+1, item[i].des);
        item[i].func();
        printf("...... done.\n\n");
        sleep(1);
    }
    return 0;
}